package live.sidian.database.synchronizer.infrastructure.sql;

import live.sidian.database.synchronizer.infrastructure.exception.SimpleSQLException;
import live.sidian.database.synchronizer.infrastructure.function.Consumer;
import live.sidian.database.synchronizer.infrastructure.utils.Assert;

import java.sql.*;
import java.util.*;
import java.util.stream.Collectors;

/**
 * @author sidian
 * @date 2021/4/2 19:01
 */
public class SqlHelper {
    Connection connection;

    public SqlHelper(String jdbcUrl, String name, String password) {
        try {
            connection = DriverManager.getConnection(jdbcUrl, name, password);
        } catch (SQLException e) {
            throw new SimpleSQLException(e);
        }
    }

    public void close() {
        if (connection == null) {
            return;
        }
        try {
            connection.close();
        } catch (SQLException ignore) {
        }
    }

    /**
     * 查询多条记录
     */
    public List<Record> queryList(String sql) {
        try {
            return doQueryList(connection, sql);
        } catch (SQLException e) {
            throw new SimpleSQLException(e);
        }
    }

    public List<Record> queryList(String sql, Consumer<PreparedStatement> consumer) {
        try {
            return doQueryList(connection, sql, consumer);
        } catch (Exception e) {
            throw new SimpleSQLException(e);
        }
    }

    /**
     * 查询单条记录
     */
    public Record queryOne(String sql) {
        List<Record> records = queryList(sql);
        Assert.isTrue(records.size() <= 1, "存在多条记录", SimpleSQLException.class);
        return records.isEmpty() ? null : records.get(0);
    }

    public Record queryOne(String sql, Consumer<PreparedStatement> consumer) {
        List<Record> records = queryList(sql, consumer);
        Assert.isTrue(records.size() <= 1, "存在多条记录", SimpleSQLException.class);
        return records.isEmpty() ? null : records.get(0);
    }

    private List<Record> doQueryList(Connection con, String sql) throws SQLException {
        // 执行sql
        Statement stmt = con.createStatement();
        ResultSet resultSet = stmt.executeQuery(sql);
        // 转化结果
        List<Record> records = parseResult(resultSet);
        stmt.close();
        resultSet.close();
        return records;
    }

    private List<Record> doQueryList(Connection con, String sql, Consumer<PreparedStatement> consumer) throws Exception {
        PreparedStatement statement = con.prepareStatement(sql);
        consumer.accept(statement);
        ResultSet resultSet = statement.executeQuery();
        List<Record> records = parseResult(resultSet);
        statement.close();
        resultSet.close();
        return records;
    }

    private List<Record> parseResult(ResultSet resultSet) throws SQLException {
        List<Record> res = new LinkedList<>();
        ResultSetMetaData metadata = resultSet.getMetaData();
        int columnCount = metadata.getColumnCount();
        while (resultSet.next()) {
            Record record = new Record();
            for (int i = 1; i <= columnCount; i++) {
                String column = metadata.getColumnName(i);
                Object value = resultSet.getObject(i);
                record.put(column, value);
            }
            res.add(record);
        }
        return res;
    }

    /**
     * 获取列与类型的映射关系
     *
     * @param table 表名
     * @return 列与类型的映射关系
     */
    public Map<String, Integer> getColumnTypeMap(String table) {
        try {
            return doGetColumnTypeMap(table);
        } catch (SQLException e) {
            throw new SimpleSQLException(e);
        }
    }

    public Set<String> getTables() {
        return queryList("show tables").stream()
                .map(record -> record.getString("Tables_in_test"))
                .collect(Collectors.toSet());
    }

    private Map<String, Integer> doGetColumnTypeMap(String table) throws SQLException {
        Map<String, Integer> map = new HashMap<>();

        Statement statement = connection.createStatement();
        ResultSet resultSet = statement.executeQuery(String.format("select * from `%s` limit 0", table));
        ResultSetMetaData metaData = resultSet.getMetaData();
        int columnCount = metaData.getColumnCount();
        for (int i = 1; i <= columnCount; i++) {
            map.put(metaData.getColumnName(i), metaData.getColumnType(i));
        }
        statement.close();
        resultSet.close();
        return map;
    }

    /**
     * 获取主键的列
     * todo 若无主键, 则可考虑唯一索引
     *
     * @param table 表名
     * @return 主键列的集合
     */
    public Set<String> getPriKeyCols(String table) {
        // 查询主键
        Set<String> cols = queryList(String.format("show keys from `%s`", table)).stream()
                .filter(record -> record.getString("Key_name").equals("PRIMARY"))
                .map(record -> record.getString("Column_name"))
                .collect(Collectors.toSet());
        if (!cols.isEmpty()) {
            return cols;
        }
        // 所有列
        return getColumnTypeMap(table).keySet();
    }

    /**
     * 执行非查询语句
     */
    public void execute(String sql) {
        try {
            doExecute(sql);
        } catch (SQLException e) {
            throw new SimpleSQLException(e);
        }
    }

    /**
     * 执行非查询语句
     */
    public void execute(String sql, Consumer<PreparedStatement> consumer) {
        try {
            doExecute(sql, consumer);
        } catch (Exception e) {
            throw new SimpleSQLException(e);
        }
    }

    private void doExecute(String sql, Consumer<PreparedStatement> consumer) throws Exception {
        PreparedStatement statement = connection.prepareStatement(sql);
        consumer.accept(statement);
        statement.execute();
        statement.close();
    }

    private void doExecute(String sql) throws SQLException {
        Statement stmt = connection.createStatement();
        stmt.execute(sql);
        stmt.close();
    }

    /**
     * 字符串转义
     */
    public static String escapeStr(String s) {
        return s.replaceAll("'", "''");
    }

    public static class Record extends HashMap<String, Object> {
        public String getString(String colName) {
            return (String) get(colName);
        }

        public Integer getInteger(String colName) {
            return (Integer) get(colName);
        }

        public Long getLong(String colName) {
            return (Long) get(colName);
        }
    }
}
